#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sun Jun  8 15:04:27 2025

@author: fenghongli
"""

import os
import pandas as pd
import matplotlib.pyplot as plt

# 输入/输出文件夹
input_dir = 'garch_volatility'
output_dir = 'volatility_plots'
os.makedirs(output_dir, exist_ok=True)

# 设置重要事件（用于可选标注）
important_dates = {
    '2011-01-26': 'New Eight Measures (Real Estate Regulations)',
    '2015-06-12': '2015Stock Market Crash',
    '2016-01-06': 'Circuit Breaker Mechanism',
    '2018-03-27': 'China-US Trade Friction',
    '2020-01-20': 'COVID-19 Pandemic',
    '2021-09-14': 'Evergrande Debt Crisis',
    '2024-05-17': 'Property Market Regulations'
}

# 可视化样式
plt.style.use('seaborn-v0_8-darkgrid')

# 遍历每个波动率文件
for file in os.listdir(input_dir):
    if file.endswith('_volatility.csv'):
        try:
            file_path = os.path.join(input_dir, file)
            df = pd.read_csv(file_path)
            df['trade_date'] = pd.to_datetime(df['trade_date'], format='%Y%m%d', errors='coerce')

            ts_code = file.replace('_volatility.csv', '')
            plt.figure(figsize=(12, 5))
            plt.plot(df['trade_date'], df['volatility'], label='波动率', linewidth=1.5, color='blue')
            plt.title(f'{ts_code} -Volatility Evolution Chart（GARCH(1,1)）')
            plt.xlabel('Transaction Date')
            plt.ylabel('Conditional volatility')
            plt.tight_layout()

            # 添加重大事件垂直线标注
            for date_str, label in important_dates.items():
                date_obj = pd.to_datetime(date_str)
                if df['trade_date'].min() <= date_obj <= df['trade_date'].max():
                    plt.axvline(x=date_obj, color='red', linestyle='--', alpha=0.6)
                    plt.text(date_obj, plt.ylim()[1]*0.8, label, color='red', fontsize=9, rotation=90)

            # 保存图像
            output_path = os.path.join(output_dir, f'{ts_code}_volatility.png')
            plt.savefig(output_path, dpi=300)
            plt.close()
        except Exception as e:
            print(f"绘制 {file} 时出错：{e}")

print(f"所有波动率图已保存至：{output_dir}/")
